import pandas as pd
from consts import (
  FEATURE_KEYWORD,
  FEATURE_GENRE,
  FEATURE_INFO,
  FEATURE_ALL,
)
from utils import (
  log_message
)

from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate

# 获取完整的训练数据
def get_train_test_data(feature_type, random_state=42):
  data_path = "../data/llm-pretrain-data/"
  train_file_path = data_path + f"train_set_random_state_{random_state}_dataset_1k.csv"
  test_file_path = data_path + f"test_set_random_state_{random_state}_dataset_1k.csv"
  try:
    """
    sep:指定分隔符
    low_memory:如果为False，则pandas将尝试使用尽可能少的内存来解析数据。如果为True，则pandas将使用更多的内存来解析数据，但可能会更快。
    """
    train_data = pd.read_csv(filepath_or_buffer=train_file_path, sep=",", low_memory=False)
    test_data = pd.read_csv(filepath_or_buffer=test_file_path, sep=",")# 这里需要对movieId进行处理，源数据的movieId是与tmdbId对应的movieId是不同的
    train_data = train_data.rename(columns={"movieId": "originalMovieId"})
    test_data = test_data.rename(columns={"movieId": "originalMovieId"})

    train_data = train_data[["userId", "originalMovieId", "title", "rating", "tmdbId"]]
    test_data = test_data[["userId", "originalMovieId", "title", "rating", "tmdbId"]]

    if feature_type == FEATURE_KEYWORD:
      file_path = data_path + "keywords.csv"
    elif feature_type == FEATURE_GENRE:
      file_path = data_path + "genres.csv"
    elif feature_type == FEATURE_INFO:
      file_path = data_path + "infos.csv"
    elif feature_type == FEATURE_ALL:
      file_path = data_path + "all.csv"

    data = pd.read_csv(filepath_or_buffer=file_path, sep=",", low_memory=False)
    # 最终关联两者根据id拿到标题和关键词相关的训练数据
    train_data = pd.merge(train_data, data, left_on='tmdbId', right_on='movieId', how='left')
    test_data = pd.merge(test_data, data, left_on='tmdbId', right_on='movieId', how='left')
    return {
      "train_data": train_data,
      "test_data": test_data
    }
  except Exception as e:
    log_message("获取文件路径错误")
    log_message(e)

def create_prompt(feature_type, train_records, test_records):
  if feature_type == FEATURE_KEYWORD:
    system_content ="the keywords of the movie"
    test_movie_template = "{title}({keywords})"
    train_input_variables = ["title", "rating", "keywords"]
  elif feature_type == FEATURE_GENRE:
    system_content ="movie's genres"
    test_movie_template = "{title}({genres})"
    train_input_variables = ["title", "rating", "genres"]
  elif feature_type == FEATURE_INFO:
    system_content ="movie's data, including title, overview, revenue, release date, average rating, popularity, vote count and rating"
    test_movie_template = "{title}({overview})({revenue})({release_date})({vote_average})({popularity})({vote_count})"
    train_input_variables = ["title", "rating", "overview", "revenue", "release_date", "vote_average", "popularity", "vote_count"]
  elif feature_type == FEATURE_ALL:
    system_content ="movie's data, including title, keywords, genres, overview, revenue, release date, average rating, popularity, vote count and rating"
    test_movie_template = "{title}({keywords})({genres})({overview})({revenue})({release_date})({vote_average})({popularity})({vote_count})"
    train_input_variables = ["title", "keywords", "genres", "rating", "overview", "revenue", "release_date", "vote_average", "popularity", "vote_count"]

  system_prompt = f"""Based on the user's past rating history and {system_content}
  You're required to predict user ratings for the candidate movie.
  The rating ranges from 0 to 5 with increments of 0.5 allowed.
  Please provide a single NUMBER as rating WITHOUT saying anything else. Do not give reasoning.
  """

  # 在测试的时候，是不需要rating的
  test_input_variables = train_input_variables[:]
  test_input_variables.remove("rating")
  train_movie_template = test_movie_template + ":{rating}"

  example_prompt = PromptTemplate(
      input_variables=train_input_variables,
      template=train_movie_template,
  )

  prefix = "The user history are:"
  pre_example_prompt = FewShotPromptTemplate(
    examples=train_records,
    example_prompt=example_prompt,
    example_separator="\n",
    prefix=prefix,
    suffix=f"\n\nThe candidate movie is \"{test_movie_template}",
    input_variables=test_input_variables,
  )
  obj = test_records[0]
  full_prompt = pre_example_prompt.format(**obj)

  return {
    "system_prompt": system_prompt,
    "full_prompt": full_prompt
  }
